from collections import deque
from typing import Deque, List, Optional, Tuple
from chernc.treediff.gumtree.matchers.mapping_store import MappingStore
from chernc.treediff.gumtree.matchers.optimal.zs.string_metrics import StringMetrics
from chernc.treediff.gumtree.tree.tree import TreeLike


class ZsMatcher:
    def __init__(self):
        self.mappings: Optional[MappingStore] = None
        self.zs_src: Optional[ZsTree] = None
        self.zs_dst: Optional[ZsTree] = None
        self.tree_dist: Optional[List[List[float]]] = None
        self.forest_dist: Optional[List[List[float]]] = None

    def match(self, src: 'TreeLike', dst: 'TreeLike', mappings: 'MappingStore') -> 'MappingStore':
        if src is None or dst is None or mappings is None:
            raise ValueError("src, dst, and mappings must not be None")

        self.zs_src = ZsTree(src)
        self.zs_dst = ZsTree(dst)
        self.mappings = mappings
        self._match()
        return mappings

    @staticmethod
    def get_first_leaf(t: 'TreeLike') -> 'TreeLike':
        if t is None:
            raise ValueError("Tree t must not be None")

        current = t
        while not current.is_leaf():
            current = current.get_child(0)
        return current

    def compute_tree_dist(self) -> list[list[float]]:
        if self.zs_src is None or self.zs_dst is None:
            raise ValueError("zs_src and zs_dst must be initialized")

        self.tree_dist = [[0.0] * (self.zs_dst.node_count + 1) for _ in range(self.zs_src.node_count + 1)]
        self.forest_dist = [[0.0] * (self.zs_dst.node_count + 1) for _ in range(self.zs_src.node_count + 1)]

        for i in range(1, len(self.zs_src.kr)):
            for j in range(1, len(self.zs_dst.kr)):
                self._forest_dist(self.zs_src.kr[i], self.zs_dst.kr[j])

        return self.tree_dist

    def _forest_dist(self, i: int, j: int):
        if self.zs_src is None or self.zs_dst is None:
            raise ValueError("zs_src and zs_dst must be initialized")
        
        if self.forest_dist is None:
            raise ValueError("forest_dist must be initialized")
        
        if self.tree_dist is None:
            raise ValueError("tree_dist must be initialized")

        self.forest_dist[self.zs_src.lld(i) - 1][self.zs_dst.lld(j) - 1] = 0
        for di in range(self.zs_src.lld(i), i + 1):
            cost_del = self.get_deletion_cost(self.zs_src.tree(di))
            self.forest_dist[di][self.zs_dst.lld(j) - 1] = self.forest_dist[di - 1][self.zs_dst.lld(j) - 1] + cost_del
            for dj in range(self.zs_dst.lld(j), j + 1):
                cost_ins = self.get_insertion_cost(self.zs_dst.tree(dj))
                self.forest_dist[self.zs_src.lld(i) - 1][dj] = self.forest_dist[self.zs_src.lld(i) - 1][dj - 1] + cost_ins

                if self.zs_src.lld(di) == self.zs_src.lld(i) and self.zs_dst.lld(dj) == self.zs_dst.lld(j):
                    cost_upd = self.get_update_cost(self.zs_src.tree(di), self.zs_dst.tree(dj))
                    self.forest_dist[di][dj] = min(
                        self.forest_dist[di - 1][dj] + cost_del,
                        self.forest_dist[di][dj - 1] + cost_ins,
                        self.forest_dist[di - 1][dj - 1] + cost_upd
                    )
                    self.tree_dist[di][dj] = self.forest_dist[di][dj]
                else:
                    self.forest_dist[di][dj] = min(
                        self.forest_dist[di - 1][dj] + cost_del,
                        self.forest_dist[di][dj - 1] + cost_ins,
                        self.forest_dist[self.zs_src.lld(di) - 1][self.zs_dst.lld(dj) - 1] + self.tree_dist[di][dj]
                    )

    def _match(self):
        if self.zs_src is None or self.zs_dst is None:
            raise ValueError("zs_src and zs_dst must be initialized")

        self.compute_tree_dist()

        root_node_pair = True
        tree_pairs: Deque[Tuple[int, int]] = deque()

        # push the pair of trees (ted1, ted2) to stack
        tree_pairs.appendleft((self.zs_src.node_count, self.zs_dst.node_count))

        while tree_pairs:
            tree_pair = tree_pairs.popleft()
            last_row, last_col = tree_pair

            # compute forest distance matrix
            if not root_node_pair:
                self._forest_dist(last_row, last_col)

            root_node_pair = False

            # compute mapping for current forest distance matrix
            first_row = self.zs_src.lld(last_row) - 1
            first_col = self.zs_dst.lld(last_col) - 1

            row, col = last_row, last_col

            while row > first_row or col > first_col:
                if row > first_row and self.forest_dist[row - 1][col] + 1.0 == self.forest_dist[row][col]:
                    # node with postorderID row is deleted from ted1
                    row -= 1
                elif col > first_col and self.forest_dist[row][col - 1] + 1.0 == self.forest_dist[row][col]:
                    # node with postorderID col is inserted into ted2
                    col -= 1
                else:
                    # node with postorderID row in ted1 is renamed to node col in ted2
                    if self.zs_src.lld(row) - 1 == self.zs_src.lld(last_row) - 1 and self.zs_dst.lld(col) - 1 == self.zs_dst.lld(last_col) - 1:
                        # if both subforests are trees, map nodes
                        t_src = self.zs_src.tree(row)
                        t_dst = self.zs_dst.tree(col)
                        if t_src.get_type() == t_dst.get_type():
                            self.mappings.add_mapping(t_src, t_dst)
                        else:
                            raise RuntimeError("Should not map incompatible nodes.")
                        row -= 1
                        col -= 1
                    else:
                        # pop subtree pair
                        tree_pairs.appendleft((row, col))
                        # continue with forest to the left of the popped subtree pair
                        row = self.zs_src.lld(row) - 1
                        col = self.zs_dst.lld(col) - 1

    def get_deletion_cost(self, n: 'TreeLike') -> float:
        if n is None:
            raise ValueError("Tree n must not be None")

        return 1.0

    def get_insertion_cost(self, n: 'TreeLike') -> float:
        if n is None:
            raise ValueError("Tree n must not be None")

        return 1.0

    def get_update_cost(self, n1: 'TreeLike', n2: 'TreeLike') -> float:
        if n1 is None or n2 is None:
            raise ValueError("Trees n1 and n2 must not be None")

        if n1.get_type() == n2.get_type():
            if n1.get_label() == "" or n2.get_label() == "":
                return 1.0
            else:
                def tokenizer(s: str):
                    s = s.replace("_", " ").lower()
                    return s.split()
                return 1.0 - StringMetrics.q_grams_distance(tokenizer=tokenizer)(n1.get_label(), n2.get_label())
        else:
            return float('inf')


class ZsTree:
    def __init__(self, t: 'TreeLike'):
        if t is None:
            raise ValueError("Tree t must not be None")

        self.node_count: int = t.get_metrics().size
        self.leaf_count: int = 0
        self.llds: List[int] = [0] * self.node_count
        self.labels: List[Optional[TreeLike]] = [None] * self.node_count
        self.kr: Optional[List[int]] = None

        idx = 1
        tmp_data = {}
        for n in t.post_order():
            tmp_data[n] = idx
            self.set_itree(idx, n)
            self.set_lld(idx, tmp_data[ZsMatcher.get_first_leaf(n)])
            if n.is_leaf():
                self.leaf_count += 1
            idx += 1

        self.set_key_roots()

    def set_itree(self, i: int, tree: 'TreeLike'):
        if tree is None:
            raise ValueError("Tree tree must not be None")

        self.labels[i - 1] = tree
        if self.node_count < i:
            self.node_count = i

    def set_lld(self, i: int, lld: int):
        if lld is None:
            raise ValueError("lld must not be None")

        self.llds[i - 1] = lld - 1
        if self.node_count < i:
            self.node_count = i

    def is_leaf(self, i: int) -> bool:
        return self.lld(i) == i

    def lld(self, i: int) -> int:
        return self.llds[i - 1] + 1

    def tree(self, i: int) -> 'TreeLike':
        return self.labels[i - 1]

    def set_key_roots(self):
        self.kr = [0] * (self.leaf_count + 1)
        visited = [False] * (self.node_count + 1)
        k = len(self.kr) - 1
        for i in range(self.node_count, 0, -1):
            if not visited[self.lld(i)]:
                self.kr[k] = i
                visited[self.lld(i)] = True
                k -= 1